import os
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter
import one_hot


# 图片形状初始化
class my_dataset(Dataset):
    def __init__(self, root_dir):
        super(my_dataset, self).__init__()
        self.image_path = [os.path.join(root_dir, image_name) for image_name in os.listdir(root_dir)]
        self.transforms = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Resize((60, 160)),  # 设置读取图片的大小
                transforms.Grayscale()
            ]
        )
        print(self.image_path)

    def __len__(self):
        return self.image_path.__len__()

    def __getitem__(self, index):
        image_parh = self.image_path[index]
        image = self.transforms(Image.open(image_parh))
        label = image_parh.split("/")[-1]
        label = label.split("_")[0]
        label_tensor = one_hot.text2Vec(label)
        label_tensor = label_tensor.view(1, -1)[0]  # 144 这里是固定格式？
        return image, label_tensor


if __name__ == '__main__':
    writer = SummaryWriter("logs")
    train_data = my_dataset("./datasets/train/")
    img, label = train_data[0]
    print(img.shape, label.shape)
    writer.add_image("img", img, 1)
    writer.close()
